import sqlite3
import os

class Database:
    def __init__(self, db_path='data/stock.db'):
        self.db_path = db_path
        os.makedirs(os.path.dirname(db_path), exist_ok=True)
        self.conn = sqlite3.connect(db_path)
        self.cursor = self.conn.cursor()
        self.init_tables()
    
    def init_tables(self):
        """初始化数据库表"""
        # 股票基本信息表
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS stocks (
            code TEXT PRIMARY KEY,
            name TEXT,
            industry TEXT,
            area TEXT,
            update_time TIMESTAMP
        )''')
        
        # 股票价格表
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS stock_prices (
            code TEXT,
            date DATE,
            open REAL,
            high REAL,
            low REAL,
            close REAL,
            volume INTEGER,
            PRIMARY KEY (code, date)
        )''')
        
        # 技术指标表
        self.cursor.execute('''
        CREATE TABLE IF NOT EXISTS technical_indicators (
            code TEXT,
            date DATE,
            ma5 REAL,
            ma10 REAL,
            ma20 REAL,
            macd REAL,
            kdj_k REAL,
            kdj_d REAL,
            kdj_j REAL,
            PRIMARY KEY (code, date)
        )''')
        
        self.conn.commit()
    
    def close(self):
        """关闭数据库连接"""
        self.conn.close() 